package top.atcra.rpc.protection.health;

import cn.hutool.core.lang.Pair;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFutureListener;
import lombok.extern.slf4j.Slf4j;
import top.atcra.rpc.RPCBootstrap;
import top.atcra.rpc.compress.CompressorType;
import top.atcra.rpc.config.AppConfig;
import top.atcra.rpc.config.ConfigLoader;
import top.atcra.rpc.transport.message.RPCRequest;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.Map;
import java.util.TimerTask;
import java.util.concurrent.*;

/**
 * 服务健康检测器，该类使用一个定时任务来周期检测所有服务的可用性。
 *
 * <p>考虑使用Netty自带的闲时检测机制替代。</p>
 *
 *
 * @author atcra
 * @date 2025-3-25
 *
 */
@Slf4j
public class ServiceHealthDetector {

    // 心跳检测的线程池，其中的线程设置为守护线程
    protected static final ScheduledExecutorService executor =
            Executors.newSingleThreadScheduledExecutor(r -> {
                Thread thread = new Thread(r);
                thread.setDaemon(true); // 设置为守护线程
                return thread;
            });
    // 待检测的服务地址列表
    protected final Map<InetSocketAddress, Channel> serviceHosts = new ConcurrentHashMap<>();

    protected static DetectTask detectTask;

    protected final RPCBootstrap instance;

    protected final Map<InetSocketAddress, Channel> channelCache;

    public ServiceHealthDetector(RPCBootstrap instance) {
        this.instance = instance;
        channelCache = instance.getChannelCache();
    }

    /**
     * 周期性检查服务可用性
     *
     * @param serviceName 服务名称
     * @param group       服务分组，dev,idGeneratorTest,prod等
     */
    public void checkHealth(String serviceName, String group) {
        var serviceRegistry = instance.getServiceRegistry();
        var addresses = serviceRegistry.lookup(serviceName, group);
        checkHealth(addresses);
    }

    public void checkHealth(List<InetSocketAddress> addresses) {
        for (var addr : addresses) {
            try {
                if (!channelCache.containsKey(addr)) {
                    var channel = instance.createChannel(addr);
                    channelCache.put(addr, channel);
                    serviceHosts.put(addr, channel);
                } else {
                    serviceHosts.put(addr, channelCache.get(addr));
                }
            } catch (Exception e) {
                log.error("连接失败，服务{}不可用", addr, e);
                serviceHosts.remove(addr);
            }
        }
        synchronized (ServiceHealthDetector.class) {
            if (detectTask == null) {
                detectTask = new DetectTask();
                executor.scheduleAtFixedRate(detectTask, 0, ConfigLoader.getConfig().getHealthCheckInterval_MS(), TimeUnit.MILLISECONDS);
            }
        }
    }


    protected class DetectTask extends TimerTask {
        protected static final int MAX_RETRY = 2;
        protected AppConfig cfg;


        @Override
        public void run() {
            cfg = ConfigLoader.getConfig();
            RPCBootstrap.CHANNEL_TTL_CACHE.clear();
            serviceHosts.forEach((addr, channel) -> {
                long ttl = -1;
                for (int i = 0; i < MAX_RETRY; ++i) {
                    long begin = System.currentTimeMillis(), end;
                    var callFuture = sendHeartbeat(channel);
                    try {
                        callFuture.get(cfg.getHealthCheckTimeout_MS(), TimeUnit.MILLISECONDS);
                        end = System.currentTimeMillis();
                        ttl = end - begin;
                        break;
                    } catch (ExecutionException | InterruptedException | TimeoutException e) {
                        log.warn("服务{}不可用，执行第{}次重试...", addr, i + 1, e);
                    }
                }
                updateServiceTTL(channel, ttl);
                if (ttl == -1) {
                    channelCache.remove(addr);
                    log.warn("服务{}不可用，已从缓存中移除", addr);
                } else {
                    if (!channelCache.containsKey(addr)) {
                        channelCache.put(addr, channel);
                    }
                    log.info("服务{}可用，向缓存中加入", addr);
                }
                if (log.isDebugEnabled()) {
                    log.debug("服务{}的响应时间为{}ms", addr, ttl);
                }

            });
        }
    }


    /**
     * 发送心跳包
     *
     * @param channel 服务连接
     * @return 响应Future对象
     */
    public static CompletableFuture<Object> sendHeartbeat(Channel channel) {
        RPCRequest rpcRequest = RPCRequest.builder()
                .requestId(channel.attr(RPCBootstrap.REQUEST_ID_GENERATOR_KEY).get().getId2())
                .serializeType(ConfigLoader.getConfig().getSerializer())
                .compressType(CompressorType.none.getId())
                .requestType(RPCRequest.RequestType.HEALTH_DETECT.getId())
                .timestamp(System.currentTimeMillis())
                .build();
        CompletableFuture<Object> callFuture = new CompletableFuture<>();
        channel.attr(RPCBootstrap.PENDING_CALLS_KEY).get().put(rpcRequest.getRequestId(), Pair.of(callFuture, rpcRequest.getTimestamp()));
        channel.writeAndFlush(rpcRequest).addListener((ChannelFutureListener) writtenPromise -> {
            if (!writtenPromise.isSuccess()) {
                callFuture.completeExceptionally(writtenPromise.cause());
            }
        });
        return callFuture;
    }

    public static void updateServiceTTL(Channel channel, long ttl) {
        RPCBootstrap.CHANNEL_TTL_CACHE.put(ttl, channel);
    }
}
